import torch
import torch.onnx
from model import BurstBotRGCN

# Initialize the model with the correct number of features from the training data
num_users = 8694
inters = 50
num_properties = 12  # Use 12 numerical features as in the training
cat_properties = 14  # Set to 14 to match training data (as the error suggests)
embedding_dimension1 = 16
embedding_dimension2 = 16
embedding_dimension = 32
dropout1 = 0.3
dropout2 = 0.3

# Initialize model with the correct feature sizes
model = BurstBotRGCN(num_users, inters, num_properties, cat_properties,
                     embedding_dimension1=embedding_dimension1, embedding_dimension2=embedding_dimension2,
                     dropout1=dropout1, num_prop_size=5, cat_prop_size=3, # Ensure it is set to 14 as in training
                     embedding_dimension=embedding_dimension,
                     dropout2=dropout2).to('cuda')

# Load trained weights
model.load_state_dict(torch.load('best_state_dict.pt', map_location='cuda'), strict=False)
model.eval()

# Create sample inputs based on the shapes used during training
des = torch.randn(229580, 768)  # Description input
tweet = torch.randn(229580, 768)  # Tweet input
num_prop = torch.randn(229580, 5)  # Numerical feature input (12 features)
cat_prop = torch.randn(229580, 3)  # Categorical feature input (14 features to match training)
burst_num_tensor = torch.randn(229580, num_properties)  # Burst graph numerical input (12 features)
burst_cat_tensor = torch.randn(229580, cat_properties)  # Burst graph categorical input (14 features)
tweet_range_list = torch.randint(0, 229580, (8694,))  # Tweet range list per user
re_index = torch.randint(0, 8694, (229580,))  # Reindex list
edge_index_burst = torch.randint(0, 229580, (2, 10)).long()  # Burst graph edges
edge_index_rgcn = torch.randint(0, 229580, (2, 10)).long()  # RGCN graph edges
edge_type = torch.randint(0, 2, (10,)).long()  # RGCN edge types

# Export the model to ONNX
torch.onnx.export(model,
                  (burst_num_tensor, burst_cat_tensor, tweet_range_list, edge_index_burst, re_index,
                   des, tweet, num_prop, cat_prop, edge_index_rgcn, edge_type),
                  'burstbotrgcn.onnx',
                  input_names=['burst_num_tensor', 'burst_cat_tensor', 'tweet_range_list', 'edge_index_burst', 're_index',
                               'des', 'tweet', 'num_prop', 'cat_prop', 'edge_index_rgcn', 'edge_type'],
                  output_names=['output'],
                  dynamic_axes={
                      'burst_num_tensor': {0: 'batch_size'},
                      'burst_cat_tensor': {0: 'batch_size'},
                      'des': {0: 'batch_size'},
                      'tweet': {0: 'batch_size'},
                      'num_prop': {0: 'batch_size'},
                      'cat_prop': {0: 'batch_size'},
                      'edge_index_burst': {1: 'num_edges_burst'},
                      'edge_index_rgcn': {1: 'num_edges_rgcn'},
                      'edge_type': {0: 'num_edges_rgcn'},
                      'output': {0: 'batch_size'}
                  })
